'''
Created on Dec 9, 2009

@author: mkiyer
'''
import unittest
import veggie.io.wiggle as wiggle
import os

def compare_wiggles(iter1, iter2):
    for line1, line2 in zip(iter1, iter2):
        print line1.strip(), line2
        if line1.strip() != line2:
            print line1.strip(), line2.strip()
            return False
    return True

def test_wiggles(correct_string, arrays):
    # make a temp file
    mywig = 'tmp.wig'
    # write wig
    wigglewriter = wiggle.WiggleFileWriter(mywig, compress=False)
    chrom = 'chr1'
    start = 0
    for arr in arrays:
        wigglewriter.write_fixed_step(chrom, start, start + len(arr), arr)    
        start += len(arr)
    # compare
    wigglewriter.close()
    mywig = open(mywig, 'r')
    return compare_wiggles(mywig, correct_string.split('\n'))

class Test(unittest.TestCase):

    def testExample1(self):
        arrays = [[0, 0, 0, 0, 1, 1, 2, 2, 0, 0],
                  [0, 0, 0, 0, 1, 1, 2, 2, 0, 0]]
        example1 = \
'''fixedStep chrom=chr1 start=5 step=1
1
1
2
2
fixedStep chrom=chr1 start=15 step=1
1
1
2
2
'''
        self.assertTrue(test_wiggles(example1, arrays))

    def testExample2(self):
        arrays = [[0, 0, 0, 0, 1, 1, 2, 2, 2, 3],
                  [3, 4, 0, 0, 1, 1, 2, 2, 0, 0]]
        example2 = \
'''fixedStep chrom=chr1 start=5 step=1
1
1
2
2
2
3
3
4
fixedStep chrom=chr1 start=15 step=1
1
1
2
2
'''
        self.assertTrue(test_wiggles(example2, arrays))


    def compare_variable_step(self, fhd, correct_header, correct_lines):
        self.assertEqual(fhd.readline().strip(), correct_header)        
        for line, correct_line in zip(fhd, correct_lines):
            fields = line.strip().split('\t')
            pos, val = int(fields[0]), float(fields[1])
            self.assertEqual(pos, correct_line[0])
            self.assertEqual(val, correct_line[1])

    def test_write_span(self):
        arrays = [[0, 0, 1, 1, 2, 2, 3, 3, 4, 4],
                  [5, 5, 6, 6, 7, 7, 8, 8, 9, 9],
                  [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
                  [0,0,0,0,0,0,0,0,0,0],
                  [0,0,0,0,0,0,0,0,0,0],
                  [100, 100, 100, 100, 100, 100, 100, 100, 100, 100]]
        chrom = 'chr1'
        start = 0
        span = 10
        mywig = 'tmp.wig'
        wigglewriter = wiggle.WiggleFileWriter(mywig, compress=False, span=span)
        for arr in arrays:
            wigglewriter.write_span(chrom, start, start + len(arr), arr)    
            start += len(arr)
        wigglewriter.close()
        header = 'variableStep chrom=chr1 span=%s' % (span)
        lines = [(1, 2), (11, 7), (21, 14.5), (51, 100)]
        self.compare_variable_step(open(mywig), header, lines)
        os.remove(mywig)
        del wigglewriter        
        # start at different spot
        span = 10
        start = 16 
        wigglewriter = wiggle.WiggleFileWriter('tmp.wig', compress=False, span=span)
        for arr in arrays:
            wigglewriter.write_span(chrom, start, start + len(arr), arr)    
            start += len(arr)
        wigglewriter.close()        
        header = 'variableStep chrom=chr1 span=%s' % (span)
        lines = [(11, 0.2), (21, 4), (31, 9.4), (41, 9.9), (61, 40), (71, 60)]        
        self.compare_variable_step(open(mywig), header, lines)
        os.remove(mywig)
        del wigglewriter
        

    def test_reader(self):
        import numpy as np
        lines = ['fixedStep chrom=chr1 start=1 step=1\n', '0\n', '1\n', '2\n', '3\n',
                 'fixedStep chrom=chr10 start=41 step=10 span=2\n', '5\n', '10\n',
                 'fixedStep chrom=chr1 start=1001 step=1\n', '5\n', '10\n',
                 'fixedStep chrom=chr1 start=1501 step=1\n', '20\n']                 
        result_chunks = [('chr1', 0, 4, [0, 1, 2, 3]),
                         ('chr10', 40, 50, [5, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]),
                         ('chr10', 50, 60, [10, 10, 0, 0, 0, 0, 0, 0, 0, 0]),
                         ('chr1', 1000, 1002, [5, 10]),
                         ('chr1', 1500, 1501, [20])]
        # read it back
        wigglereader = wiggle.WiggleFileReader(lines, dtype=np.int)
        print
        for chunk in wigglereader.iterchunks(10):
            start, end, arr = chunk
            print wigglereader.chrom, chunk
            #self.assertEqual((wigglereader.chrom, start, end), correct_chunk[0:3])
            
        

        

if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    unittest.main()